#include "DenseGradeWrapper.h"

DenseGradeWrapper::DenseGradeWrapper(QString xmlPath, QObject *parent)
    : QObject{parent}
{
    qDebug() << ov::get_openvino_version().buildNumber << " " << ov::get_openvino_version().description;
    qDebug() << "Start To Compile CPU Model";
    ov::Core core;
    this->xmlModel = core.read_model(xmlPath.toStdString());
    printInputAndOutputsInfo(*this->xmlModel);
    double tS, tE;
    tS = (double) clock();
    this->compiledModel = core.compile_model(this->xmlModel, "CPU");
    tE = (double) clock();
    qDebug() << "Compile Done, cost " << (tE - tS) / (CLOCKS_PER_SEC) << " s";
}

DenseInferResult DenseGradeWrapper::infer(QString imageSrc){
    double tS, tE;

    tS = (double) clock();
    cv::Mat transDim = this->preprocess(imageSrc);
    tE = (double) clock();
    qDebug() << "Preprocess Done, cost " << (tE - tS) / (CLOCKS_PER_SEC) << " s";

    tS = (double) clock();
    QVector<float> buffer = this->_infer(transDim);
    tE = (double) clock();
    qDebug() << "Infer Done, cost " << (tE - tS) / (CLOCKS_PER_SEC) << " s";

    tS = (double) clock();
    this->softmax(buffer);
    DenseInferResult res = this->postProcess(buffer);
    tE = (double) clock();
    qDebug() << "Postprocess Done, cost " << (tE - tS) / (CLOCKS_PER_SEC) << " s";

    return res;
}

QVector<float> DenseGradeWrapper::_infer(cv::Mat transDim){
    ov::InferRequest inferRequest = this->compiledModel.create_infer_request();
    ov::Shape inputShape = {1, 3, 224, 224};
    ov::Tensor inputTensor = ov::Tensor(ov::element::f32, inputShape, transDim.data);
    inferRequest.set_input_tensor(inputTensor);

    // async
    inferRequest.start_async();
    inferRequest.wait();

    const ov::Tensor& outputTensor = inferRequest.get_output_tensor();
    const float *outputBuffer = outputTensor.data<const float>();
    return QVector<float>(outputBuffer, (outputBuffer + outputTensor.get_size()));
}

void DenseGradeWrapper::softmax(QVector<float> &buffer){
    float denominator = 0.0;

    for(size_t i = 0; i < buffer.size(); i++){
        buffer[i] = std::exp(buffer[i]);
        denominator += buffer[i];
    }

    for(size_t i = 0; i < buffer.size(); i++){
        buffer[i] /= denominator;
    }
}

DenseInferResult DenseGradeWrapper::postProcess(const QVector<float> &buffer){
    int index = buffer.indexOf(*std::max_element(buffer.begin(), buffer.end()));
    DenseInferResult result;
    result.index = index;
    std::copy(buffer.begin(), buffer.end(), result.pie);
    return result;
}

cv::Mat DenseGradeWrapper::preprocess(QString imageSrc){
    cv::Mat cropMat = this->circleCrop(imageSrc);
    std::vector<cv::Mat> batch_mat;
    batch_mat.push_back(cropMat);
    return cv::dnn::blobFromImages(batch_mat, 1.0 / 255.0, cv::Size(),
                                   cv::Scalar(), true);
}


cv::Mat DenseGradeWrapper::circleCrop(QString imageSrc){
    cv::Mat cropMask = NetworkTool::ToolCropImageFromMask(imageSrc);
    return NetworkTool::ToolCropWithGaussian(cropMask);
}
